package com.template.tree;

/**
 * AVL 是一种自平衡二叉搜索树，它通过旋转操作来保持树的平衡性。在插入或删除一个节点时，AVL 树会检查是否破坏了平衡性，并通过四种旋转操作将树重新平衡。
 *
 * 在添加一个节点时，先按照二叉搜索树的规则（小于当前节点则进入左子树，大于当前节点则进入右子树）找到合适的插入位置，然后根据新插入节点的位置和父节点以及祖先节点的高度差，判断是否需要进行旋转操作。具体的旋转操作包括单旋转（左旋或右旋）和双旋转（先左后右或先右后左），旋转时要注意更新节点高度。
 * @author shenb
 * @date 2023-03-28 14:01
 */
public class AVLTree {
    private Node root;

    // 定义AVL树的节点类
    private static class Node {
        int val, height; // 节点值和高度
        Node left, right; // 左子节点和右子节点

        Node(int val) {
            this.val = val;
            height = 1;
        }
    }

    // 获取节点的高度，如果节点为空则返回0
    private int height(Node node) {
        return node == null ? 0 : node.height;
    }

    // 计算节点的平衡因子，即左右子树高度差
    private int balanceFactor(Node node) {
        return node == null ? 0 : height(node.left) - height(node.right);
    }

    // 更新节点高度
    private void updateHeight(Node node) {
        node.height = 1 + Math.max(height(node.left), height(node.right));
    }

    // 左旋操作，使得右子树成为新的根节点
    private Node rotateLeft(Node node) {
        Node right = node.right;
        Node rightLeft = right.left;
        right.left = node;
        node.right = rightLeft;
        updateHeight(node);
        updateHeight(right);
        return right;
    }

    // 右旋操作，使得左子树成为新的根节点
    private Node rotateRight(Node node) {
        Node left = node.left;
        Node leftRight = left.right;
        left.right = node;
        node.left = leftRight;
        updateHeight(node);
        updateHeight(left);
        return left;
    }

    // 重新平衡节点，通过旋转操作调整子树的高度差
    private Node rebalance(Node node) {
        updateHeight(node); // 更新节点高度
        int balanceFactor = balanceFactor(node); // 计算平衡因子
        if (balanceFactor > 1) { // 左子树高度较高
            if (balanceFactor(node.left) < 0) { // 右子树高度较高，先进行左旋操作
                node.left = rotateLeft(node.left);
            }
            return rotateRight(node); // 再进行右旋操作
        } else if (balanceFactor < -1) { // 右子树高度较高
            if (balanceFactor(node.right) > 0) { // 左子树高度较高，先进行右旋操作
                node.right = rotateRight(node.right);
            }
            return rotateLeft(node); // 再进行左旋操作
        } else {
            return node;
        }
    }

    // 插入节点
    public void insert(int val) {
        root = insert(root, val);
    }

    private Node insert(Node node, int val) {
        if (node == null) {
            return new Node(val); // 如果节点为空，则创建新节点
        }
        if (val < node.val) { // 如果插入值小于当前节点值，则插入到左子树中
            node.left = insert(node.left, val);
        } else { // 否则插入到右子树中
            node.right = insert(node.right, val);
        }
        return rebalance(node); // 插入完成后重新平衡节点
    }

    // 删除节点
    public void delete(int val) {
        root = delete(root, val);
    }

    private Node delete(Node node, int val) {
        if (node == null) {
            return null; // 如果节点为空，则返回null
        }
        if (val == node.val) { // 如果需要删除的节点为当前节点
            if (node.left == null && node.right == null) { // 如果没有子节点，则直接删除
                return null;
            }
            if (node.left == null) { // 如果只有右子节点，则将其替换为当前节点
                return node.right;
            }
            if (node.right == null) { // 如果只有左子节点，则将其替换为当前节点
                return node.left;
            }
            Node minNode = findMin(node.right); // 如果有两个子节点，则找到右子树中最小的节点
            node.val = minNode.val; // 将最小的节点值赋给当前节点
            node.right = delete(node.right, node.val); // 删除右子树中的最小节点
        } else if (val < node.val) { // 如果需要删除的值比当前节点小，则在左子树中查找并删除
            node.left = delete(node.left, val);
        } else { // 如果需要删除的值比当前节点大，则在右子树中查找并删除
            node.right = delete(node.right, val);
        }
        return rebalance(node); // 删除完成后重新平衡节点
    }

    // 查找右子树中最小的节点
    private Node findMin(Node node) {
        while (node.left != null) {
            node = node.left;
        }
        return node;
    }


    /**
     * 查找节点是否存在
     * @param data 数据
     * @return 是否存在
     */
    public boolean contains(int data) {
        return contains(data, root);
    }

    /**
     * 内部查找方法
     * @param data 数据
     * @param node 当前节点
     * @return 是否存在
     */
    private boolean contains(int data, Node node) {
        if (node == null) {
            return false;
        }
        if (data < node.val) {
            return contains(data, node.left);
        } else if (data > node.val) {
            return contains(data, node.right);
        } else {
            return true;
        }
    }

    /**
     * 中序遍历
     */
    public String inOrder() {
        StringBuilder sb = new StringBuilder();
        inOrder(root, sb);
        return sb.toString().trim(); // 去掉最后一个空格
    }

    private void inOrder(Node node, StringBuilder sb) {
        if (node == null) {
            return;
        }
        inOrder(node.left, sb);
        sb.append(node.val).append(" ");
        inOrder(node.right, sb);
    }

    /**
     * 查找节点
     * @param data 数据
     * @return 节点
     */
    public Node search(int data) {
        return search(data, root);
    }

    /**
     * 内部查找方法
     * @param data 数据
     * @param node 当前节点
     * @return 节点
     */
    private Node search(int data, Node node) {
        if (node == null) {
            return null;
        }
        if (data < node.val) {
            return search(data, node.left);
        } else if (data > node.val) {
            return search(data, node.right);
        } else {
            return node;
        }
    }
}

